from torch.utils.data import DataLoader
import torch
import sys

sys.path.append(r"D:\HGMSim\HGMSimBackEndByFlask\HGMSim")

from HGMSimDataset import functionSimDataset
from functionSim_config import *

import json


class HGMSimModelWrapper:
    """
    提供 单对单的计算  && topk的计算
    """

    def __init__(self) -> None:
        with open("config.json", "r") as file:
            self.config = json.load(file)
        bestModelPath = self.config["model"]["bestModelPath"]
        self.model = torch.load(bestModelPath, map_location=torch.device("cpu"))
        self.optimizer = torch.optim.Adam(
            [
                {"params": self.model.parameters()},
            ],
            lr=lr,
        )

    def get_model_dataloader(self, datas, is_shuffle=False):
        """
        datas用于加载其它的输入,方便其他模型的对比
        """
        my_dataset = functionSimDataset(datas)
        dataloader = DataLoader(
            my_dataset,
            batch_size=batchSize,
            shuffle=is_shuffle,
            num_workers=0,  # 不要并发加载，能显著提升速度
            collate_fn=my_dataset.adjust_samples_to_same_dimension,
        )
        return dataloader

    def getModelScore(self, dataloader):
        """
        sampleX: 目标样本名称
        sampleList: 相似性计算样本名称列表
        计算目标样本和样本列表中每一对的相似值
        """
        self.model.eval()
        with torch.no_grad():
            res = []
            for i, (query, target1, target2) in enumerate(dataloader):
                self.optimizer.zero_grad()
                adj_x, att_x, vtype_x = target1
                adj_y, att_y, vtype_y = target2
                adj_x = adj_x.to(device)
                adj_y = adj_y.to(device)
                att_x = att_x.to(device)
                att_y = att_y.to(device)
                vtype_x = vtype_x.to(device)
                vtype_y = vtype_y.to(device)
                score1 = self.model(adj_x, att_x, vtype_x, adj_y, att_y, vtype_y)
                res.append(score1)
            ans = torch.cat((res))
        return ans.tolist()

    def calculateSim(self, datas):
        dataloader = self.get_model_dataloader(datas)
        res = self.getModelScore(dataloader)
        return res


if __name__ == "__main__":
    hGMSimModelWrapper = HGMSimModelWrapper()
    a = [
        ("02d63ae96c49f0bed65dfc64404e2c59", "6e56ab5d8d53a5c8666007530f692f40", 1),
        ("02d63ae96c49f0bed65dfc64404e2c59", "7e37cc16388e6616e704c73456b6a492", 1),
        ("02d63ae96c49f0bed65dfc64404e2c59", "29c7250265ee4419dc785931fc865b78", 1),
    ]
    res = hGMSimModelWrapper.calculateSim(a)
    print(res)
